from functools import cached_property
from pathlib import Path
from typing import Any, Iterable
from typing_extensions import Self
from pydantic import BaseModel, Field
from pysam import TabixFile, VariantFile, VariantHeader, VariantRecord, tabix_compress, tabix_index  # pylint: disable=E0611
from openanno_cnv.cnv import CNV
from openanno_cnv.transcript import Transcript


class BND(BaseModel):
    chrom: str
    pos: int
    chrom2: str
    pos2: int
    info: dict[str, Any] = {}

    @cached_property
    def pk(self) -> str:
        return f"{self.chrom}:{self.pos}-{self.chrom2}:{self.pos2}"

    @classmethod
    def from_variant_record(cls, record: VariantRecord) -> Self:
        return cls(chrom=record.chrom, pos=record.start, chrom2=record.info["CHR2"], pos2=record.info["POS2"])

    def query_tabix(self, tbx: TabixFile) -> Iterable[str]:
        try:
            return tbx.fetch(self.chrom, self.start, self.end)
        except Exception:  # pylint: disable=W0718
            return []

    def anno_gene(self, tbx: TabixFile, gene_ids: dict[str, str]):
        data = {}
        for chrom, pos in ((self.chrom, self.pos), (self.chrom2, self.pos2)):
            try:
                iter_line = tbx.fetch(chrom, pos, pos + 1)
            except Exception:  # pylint: disable=W0718
                pass
            else:
                for line in iter_line:
                    transcript = Transcript.from_genepred_line(line)
                    gene_id = gene_ids.get(transcript.gene_pk, ".")
                    detail = transcript.calc_detail(pos, pos + 1)
                    data.setdefault(transcript.gene_pk, {"gene": transcript.gene, "gene_id": gene_id, "detail": ""})
                    sep = "|" if data[transcript.gene_pk]["detail"] else ""
                    data[transcript.gene_pk]["detail"] += f"{sep}{detail}"
        if data:
            genes, gene_ids, details = list(zip(*map(lambda x: (x["gene"], x["gene_id"], x["detail"]), data.values())))
            self.info.update({"GENE": genes, "GENE_ID": gene_ids, "DETAIL": details})


class Annotator(BaseModel):
    bcf_file: Path
    vcf_file: Path
    genepred_file: Path | None = Field(default=None, description="Representative NCBI RefSeq file")
    gene_id_file: Path | None = Field(default=None, description="Gene to Entrez ID file")

    @cached_property
    def gene_ids(self) -> dict[str, str]:
        with open(self.gene_id_file) as reader:
            return dict(map(lambda x: (f"{x[0]}:{x[1]}", x[2]), map(lambda y: y.strip().split("\t"), reader)))

    @cached_property
    def annotation_header(self) -> VariantHeader:
        header = VariantHeader()
        if self.gene_id_file:
            header.info.add(id="GENE", number=".", type="String", description="Gene Symbol")
            header.info.add(id="GENE_ID", number=".", type="String", description="Gene Entrez ID")
            header.info.add(id="DETAIL", number=".", type="String", description="Gene Detail")
        return header

    @cached_property
    def genepred_tabix(self) -> TabixFile | None:
        return TabixFile(str(self.genepred_file)) if self.genepred_file else None

    def run(self):
        vcf_file = self.vcf_file.with_suffix("") if self.vcf_file.suffix == ".gz" else self.vcf_file
        with VariantFile(str(self.bcf_file)) as reader, open(vcf_file, "w") as writer:
            reader.header.merge(self.annotation_header)
            header = reader.header.copy()
            header.info.remove_header("PRECISE")
            header.info.remove_header("IMPRECISE")
            header.info.remove_header("SVMETHOD")
            writer.write(str(header))
            for record in filter(lambda x: x.info.get("PRECISE"), reader):
                sv_type = record.info.get("SVTYPE")
                if sv_type in ["INV", "BND"]:
                    sv = BND.from_variant_record(record) if sv_type == "BND" else CNV.from_variant_record(record)
                    record.alts = (f"<{sv_type}>",)
                    record.ref = "<DIP>"
                    if self.genepred_tabix:
                        sv.anno_gene(self.genepred_tabix, self.gene_ids)
                        record.info.pop("PRECISE")
                        record.info.pop("SVMETHOD")
                        record.info.update(sv.info)
                        writer.write(str(record))
        if self.vcf_file.suffix == ".gz":
            tabix_compress(str(vcf_file), str(self.vcf_file), force=True)
            tabix_index(str(self.vcf_file), force=True, preset="vcf")
            vcf_file.unlink(missing_ok=True)
